# vim:fileencoding=utf-8

"""下载 tracker 列表并更新 aria2 的配置文件文件。
    参数是配置文件名，可以是绝对路径，也可以相对于本脚本。
"""

from urllib import request
from pathlib import Path
import sys
import io
import zipfile

# 中国可用的 trackers 的下载地址。
TRACKERS_URL = "https://gitee.com/harvey520/www.yaozuopan.top/raw/master/blacklist.txt"
TRACKERS_FILE = "https://gitee.com/harvey520/top.yaozuopan/raw/master/web/iplist.txt.zip"
FILE_NAME = "iplist.txt"

def get_trackers_file(url):
    """先把文件下载下来，返回一个压缩文件字节流
    """
    res = None
    try:
        with request.urlopen(url) as stream:
            res = io.BytesIO(stream.read())
    except Exception:
        print("下载解析失败，请检查 url")
        sys.exit()
    return res

def unzip(t_file):
    """解压缩 zip，返回文件对象
    """
    return zipfile.ZipFile(t_file).read(FILE_NAME)

def get_tr_list(trackers_file, num):
    """从 trackers_file 文件对象中选出部分生成列表并转成配置文件中需要的文本字符串
    """
    res = trackers_file.decode('utf8')
    res = res.splitlines()
    res = res[0: num]
    res = ','.join(res)
    return res

def write_config(config_file, tr_list):
    """把新列表插入配置文件
    """
    item = 'bt-tracker=' + tr_list
    no_item = True
    f_path = Path(config_file)
    lines = f_path.read_text(encoding="utf-8").split('\n')
    for num, line in enumerate(lines):
        if 'bt-tracker=' == line.lstrip()[:11]:
            no_item = False
            lines[num] = item
            break
    if no_item:
        lines.append(item)
    f_path.write_text('\n'.join(lines), encoding="utf-8")

def main():
    """主函数
        没什么好说的
    """
    config_file = sys.argv[1] if len(sys.argv) > 1 else 'aria2.conf'
    tmp = get_trackers_file(TRACKERS_FILE)
    tmp = unzip(tmp)
    tr_list = get_tr_list(tmp, 20)
    write_config(config_file, tr_list)

if __name__ == "__main__":
    main()
